import asyncio
import websockets
import json

import os
import sys

import colorama as cm

import wss

# 初始化colorama库
cm.init()

config = {
    "allowed_origins": ["http://localhost:8080", "http://127.0.0.1:8080"],
    "export_path": "./export",
    "port": 6060,
    # 程序根目录
    "dir_root": "",
}


def save_json(path, data):
    dir = os.path.dirname(path)
    if not os.path.exists(dir):
        os.makedirs(dir)
    with open(path, "w") as file:
        file.write(json.dumps(data, indent=2))


def init_conf(conf_path, conf=None):
    data = ""

    if conf:
        data = json.loads(conf)
    elif os.path.exists(conf_path):
        with open(conf_path, "r") as file:
            contents = file.read()

            data = json.loads(contents)
    else:
        os.makedirs(os.path.dirname(conf_path))
        save_json(conf_path, config)
        return

    for key, value in data.items():
        config[key] = value

        save_json(conf_path, config)


# 处理 WebSocket 连接的函数
async def handle_websocket(websocket, path):
    while True:

        origin = websocket.request_headers.get("Origin", "")

        if config["allowed_origins"] == "*" or origin in config["allowed_origins"]:
            msg = await websocket.recv()

            data = json.loads(msg)

            match data["req"]:
                case "close":
                    send_msg = json.dumps({"msg": "ok", "code": 0})
                    # 在这里断开websocket链接
                    await websocket.close()  # 断开WebSocket连接
                    print("WebSocket 连接已关闭")
                    return  # 退出处理程序
                case _:
                    send_msg = wss.handle_ws_msg(msg, config)

            response = send_msg
        else:
            response = json.dumps({"msg": "拒绝请求", "code": 403})

        # 回复客户端
        await websocket.send(response)


# 启动 WebSocket 服务器
async def start_server():
    server = await websockets.serve(
        handle_websocket,
        "localhost",
        config["port"],
        extra_headers={
            "Access-Control-Allow-Origin": "*",
            "Access-Control-Allow-Methods": "GET, POST, PUT, DELETE, OPTIONS",
            "Access-Control-Allow-Headers": "Origin, Accept, Content-Type, X-Requested-With, X-CSRF-Token",
        },
    )

    print(f"WebSocket 服务已在端口 {server.sockets[0].getsockname()[1]} 上运行")

    await server.wait_closed()


# 配置文件
config_path = "config/init.json"


def run_server():
    line = f"{cm.Fore.WHITE}\n--------------------------\n"

    # print(f'{line}\033[31m 关闭该窗体以停止 websocket 服务\033[0m{line}\n')
    print(
        f"{line}{cm.Fore.RED} 关闭该窗体以停止 websocket 服务{cm.Fore.WHITE}{line}{cm.Fore.WHITE}\n"
    )

    print(f"配置文件位于: {cm.Fore.CYAN}{os.path.abspath(config_path)}")
    # print(f'\033[31mconfig\033[0m = \033[32m{json.dumps(config, indent=4)}\033[0m')
    print(
        f"\n{cm.Fore.RED}config{cm.Fore.WHITE} = {cm.Fore.GREEN}{json.dumps(config, indent=4)}{cm.Fore.WHITE}"
    )

    print(line)

    # 运行 WebSocket 服务器
    asyncio.run(start_server())


isInit = False

if len(sys.argv) < 2:
    # print(f'未传入参数，使用本地配置: \033[34m{os.path.abspath(config_path)}\033[0m')
    print(f"未传入参数，使用本地配置: {cm.Fore.BLUE}{os.path.abspath(config_path)}")

    init_conf(config_path)
else:
    data = json.loads(sys.argv[1])
    print(sys.argv[1])

    if "init" in data:
        del data["init"]
        isInit = True

    config_path = os.path.join(data["dir_root"], config_path)
    print(config_path)

    for key, value in data.items():
        config[key] = value

    save_json(config_path, config)

if isInit:
    sys.exit(config["port"])
else:
    os.system(f'title ws://localhost:{config["port"]}')
    run_server()
